package com.example.deepai.service;

import io.milvus.client.MilvusServiceClient;
import io.milvus.param.dml.InsertParam;
import io.milvus.param.dml.SearchParam;
import io.milvus.grpc.MutationResult;
import io.milvus.grpc.SearchResults;
import io.milvus.param.MetricType;
import io.milvus.param.R;
import io.milvus.param.collection.LoadCollectionParam;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;

import java.util.Arrays;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.stream.Collectors;

@Service
@Slf4j
public class KnowledgeGraphService {
    private final MilvusServiceClient milvusClient;
    private static final String COLLECTION_NAME = "ui_analysis";
    private static final int VECTOR_DIMENSION = 512; // 确保与集合创建时的维度一致
    
    public KnowledgeGraphService(MilvusServiceClient milvusClient) {
        this.milvusClient = milvusClient;
        // 初始化时加载集合
        loadCollection();
    }
    
    private void loadCollection() {
        try {
            // 检查集合是否已加载
            boolean isLoaded = false;
            try {
                // 这里可以添加检查集合是否已加载的逻辑
                // 例如使用 showCollections 或 describeCollection API
                isLoaded = true; // 假设已加载，实际应该根据 API 返回判断
            } catch (Exception e) {
                log.warn("检查集合加载状态失败: {}", e.getMessage());
            }
            
            if (!isLoaded) {
                LoadCollectionParam loadParam = LoadCollectionParam.newBuilder()
                    .withCollectionName(COLLECTION_NAME)
                    .build();
                milvusClient.loadCollection(loadParam);
                log.info("Collection {} loaded successfully", COLLECTION_NAME);
                
                // 增加等待时间，确保集合完全加载
                Thread.sleep(3000);
            } else {
                log.info("Collection {} already loaded", COLLECTION_NAME);
            }
        } catch (Exception e) {
            log.warn("Failed to load collection {}: {}", COLLECTION_NAME, e.getMessage());
        }
    }
    
    public List<String> searchSimilarAnalysis(float[] queryEmbedding) {
        try {
            // 确保集合已加载
            loadCollection();
            
            // 调整向量维度
            float[] adjustedEmbedding = adjustVectorDimension(queryEmbedding, VECTOR_DIMENSION);
            
            // 将 float[] 转换为 List<Float>
            List<Float> embeddingList = new ArrayList<>(adjustedEmbedding.length);
            for (float v : adjustedEmbedding) {
                embeddingList.add(v);
            }
            
            // 重试机制
            int maxRetries = 5; // 增加重试次数
            for (int i = 0; i < maxRetries; i++) {
                try {
                    SearchParam searchParam = SearchParam.newBuilder()
                        .withCollectionName(COLLECTION_NAME)
                        .withMetricType(MetricType.L2)  // 使用欧氏距离
                        .withOutFields(Arrays.asList("type", "description"))
                        .withTopK(3)
                        .withVectors(Collections.singletonList(embeddingList))
                        .withVectorFieldName("embedding")
                        .build();
                        
                    R<SearchResults> response = milvusClient.search(searchParam);
                    if (response.getStatus() != 0) {
                        log.error("搜索分析失败 (尝试 {}/{}): {}", i+1, maxRetries, response.getMessage());
                        if (i < maxRetries - 1) {
                            // 指数退避策略，每次等待时间增加
                            Thread.sleep(1000 * (i + 1));
                            // 重新加载集合
                            loadCollection();
                            continue;
                        }
                        return Collections.emptyList();
                    }
                    
                    // 检查结果是否为空
                    if (response.getData().getResults().getFieldsDataCount() == 0) {
                        log.warn("搜索结果为空");
                        return Collections.emptyList();
                    }
                    
                    return response.getData()
                        .getResults()
                        .getFieldsDataList()
                        .stream()
                        .map(field -> field.getScalars().getStringData().getData(0))
                        .collect(Collectors.toList());
                } catch (Exception e) {
                    log.error("搜索分析时发生异常 (尝试 {}/{}): {}", i+1, maxRetries, e.getMessage());
                    if (i < maxRetries - 1) {
                        Thread.sleep(1000 * (i + 1));
                    }
                }
            }
            
            return Collections.emptyList();
        } catch (Exception e) {
            log.error("搜索分析时发生异常: {}", e.getMessage(), e);
            return Collections.emptyList();
        }
    }
    
    /**
     * 调整向量维度到目标维度
     * @param vector 原始向量
     * @param targetDimension 目标维度
     * @return 调整后的向量
     */
    private float[] adjustVectorDimension(float[] vector, int targetDimension) {
        if (vector.length == targetDimension) {
            return vector; // 维度已匹配，无需调整
        }
        
        float[] adjustedVector = new float[targetDimension];
        
        if (vector.length > targetDimension) {
            // 截断向量
            System.arraycopy(vector, 0, adjustedVector, 0, targetDimension);
            log.info("向量维度已从 {} 截断至 {}", vector.length, targetDimension);
        } else {
            // 填充向量（用0填充）
            System.arraycopy(vector, 0, adjustedVector, 0, vector.length);
            // 剩余部分默认为0
            log.info("向量维度已从 {} 填充至 {}", vector.length, targetDimension);
        }
        
        return adjustedVector;
    }
    
    /**
     * 截断字符串到指定长度
     * @param text 原始字符串
     * @param maxLength 最大长度
     * @return 截断后的字符串
     */
    private String truncateString(String text, int maxLength) {
        if (text == null) {
            return "";
        }
        if (text.length() <= maxLength) {
            return text;
        }
        return text.substring(0, maxLength);
    }
    
    public void storeAnalysis(String type, String description, float[] embedding) {
        try {
            // 截断描述文本，防止超出Milvus字段长度限制
            String truncatedType = truncateString(type, 100);
            String truncatedDescription = truncateString(description, 500);
            
            // 调整向量维度
            float[] adjustedEmbedding = adjustVectorDimension(embedding, VECTOR_DIMENSION);
            
            // 将 float[] 转换为 List<Float>
            List<Float> embeddingList = new ArrayList<>(adjustedEmbedding.length);
            for (float v : adjustedEmbedding) {
                embeddingList.add(v);
            }
            
            List<InsertParam.Field> fields = new ArrayList<>();
            
            fields.add(new InsertParam.Field("type", Collections.singletonList(truncatedType)));
            fields.add(new InsertParam.Field("description", Collections.singletonList(truncatedDescription)));
            fields.add(new InsertParam.Field("embedding", Collections.singletonList(embeddingList)));
            
            InsertParam insertParam = InsertParam.newBuilder()
                .withCollectionName(COLLECTION_NAME)
                .withFields(fields)
                .build();
            
            R<MutationResult> response = milvusClient.insert(insertParam);
            if (response.getStatus() != 0) {
                log.error("存储分析失败: {}", response.getMessage());
            } else {
                log.info("成功存储分析数据，类型: {}", truncatedType);
            }
        } catch (Exception e) {
            log.error("存储分析时发生异常: {}", e.getMessage(), e);
        }
    }
}